import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import os
from mpl_toolkits.mplot3d import Axes3D  # 导入3D绘图模块

# .csv数据文件目录
directory = "src/build/"

# 读取数据并绘图
i_values = [1, 2, 3]
n_values = [10, 40, 160]

# 处理累积弦长拟合
for i in i_values:
    fig = plt.figure(figsize=(10, 8)) # 为每个i创建一个新的图形
    if i == 1 or i == 2:
        # 绘制精确函数（仅读取一次）
        data1 = pd.read_csv(f'{directory}data_E{i}_cumulative.csv', header=None)
        plt.plot(data1[0], data1[1], 'k--', label='Exact function')

    elif i == 3:
        # 对于i=3，绘制三维数据
        ax = fig.add_subplot(111, projection='3d')  # 创建一个3D坐标轴

        # 绘制精确函数（仅读取一次）
        data1 = pd.read_csv(f'{directory}data_E{i}_cumulative.csv', header=None)
        ax.plot(data1[0], data1[1], data1[2], 'k--', label='Exact function')

    for n in n_values:
        # 读取累计弦长数据
        data = pd.read_csv(f'{directory}data_E{i}_{n}_cumulative.csv', header=None)  # 指定没有列标题
        
        if i == 1 or i == 2:
            # 绘制二维数据点
            plt.plot(data[0], data[1], label=f'n={n}')
            plt.title(f'Problem E cumulative for i={i}')

        elif i == 3:
            # 绘制3D数据
            ax.plot(data[0], data[1], data[2], label=f'n={n}')

            # 添加标签
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_zlabel('Z')
            ax.set_title(f'Problem E cumulative for i={i}')

    # 添加图例
    plt.legend(loc='best')
    plt.savefig(f'figure/E{i}_cumulative.png')
    plt.close()

# 处理等距节点拟合
for i in i_values:
    fig = plt.figure(figsize=(10, 8)) # 为每个i创建一个新的图形
    if i == 1 or i == 2:
        # 绘制精确函数（仅读取一次）
        data1 = pd.read_csv(f'{directory}data_E{i}_equidistant.csv', header=None)
        plt.plot(data1[0], data1[1], 'k--', label='Exact function')

    elif i == 3:
        # 对于i=3，绘制三维数据
        ax = fig.add_subplot(111, projection='3d')  # 创建一个3D坐标轴

        # 绘制精确函数（仅读取一次）
        data1 = pd.read_csv(f'{directory}data_E{i}_equidistant.csv', header=None)
        ax.plot(data1[0], data1[1], data1[2], 'k--', label='Exact function')

    for n in n_values:
        # 读取等距节点数据
        data = pd.read_csv(f'{directory}data_E{i}_{n}_equidistant.csv', header=None)  # 指定没有列标题
        
        if i == 1 or i == 2:
            # 绘制二维数据点
            plt.plot(data[0], data[1], label=f'n={n}')
            plt.title(f'Problem E equidistant for i={i}')

        elif i == 3:     
            # 绘制3D数据
            ax.plot(data[0], data[1], data[2], label=f'n={n}')

            # 添加标签
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.set_zlabel('Z')
            ax.set_title(f'Problem E equidistant for i={i}')

    # 添加图例
    plt.legend(loc='best')
    plt.savefig(f'figure/E{i}_equidistant.png')
    plt.close()
